{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from mpc import mpc\n",
    "from mpc.mpc import QuadCost, LinDx, GradMethods\n",
    "from mpc.env_dx import cartpole\n",
    "\n",
    "import numpy as np\n",
    "import numpy.random as npr\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import os\n",
    "import io\n",
    "import base64\n",
    "import tempfile\n",
    "from IPython.display import HTML\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/100 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tmp dir: /var/folders/6c/cyr3grs541gf55mv77xsplnw0000gn/T/tmpx36g3yql\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [01:28<00:00,  1.93it/s]\n"
     ]
    }
   ],
   "source": [
    "dx = cartpole.CartpoleDx()\n",
    "\n",
    "n_batch, T, mpc_T = 8, 100, 25\n",
    "\n",
    "def uniform(shape, low, high):\n",
    "    r = high-low\n",
    "    return torch.rand(shape)*r+low\n",
    "\n",
    "torch.manual_seed(0)\n",
    "th = uniform(n_batch, -2*np.pi, 2*np.pi)\n",
    "thdot = uniform(n_batch, -.5, .5)\n",
    "x = uniform(n_batch, -0.5, 0.5)\n",
    "xdot = uniform(n_batch, -0.5, 0.5)\n",
    "xinit = torch.stack((x, xdot, torch.cos(th), torch.sin(th), thdot), dim=1)\n",
    "\n",
    "x = xinit\n",
    "u_init = None\n",
    "\n",
    "q, p = dx.get_true_obj()\n",
    "Q = torch.diag(q).unsqueeze(0).unsqueeze(0).repeat(\n",
    "    mpc_T, n_batch, 1, 1\n",
    ")\n",
    "p = p.unsqueeze(0).repeat(mpc_T, n_batch, 1)\n",
    "\n",
    "t_dir = tempfile.mkdtemp()\n",
    "print('Tmp dir: {}'.format(t_dir))\n",
    "\n",
    "action_history = []\n",
    "for t in tqdm(range(T)):\n",
    "    nominal_states, nominal_actions, nominal_objs = mpc.MPC(\n",
    "        dx.n_state, dx.n_ctrl, mpc_T,\n",
    "        u_init=u_init,\n",
    "        u_lower=dx.lower, u_upper=dx.upper,\n",
    "        lqr_iter=50,\n",
    "        verbose=0,\n",
    "        exit_unconverged=False,\n",
    "        detach_unconverged=False,\n",
    "        linesearch_decay=dx.linesearch_decay,\n",
    "        max_linesearch_iter=dx.max_linesearch_iter,\n",
    "        grad_method=GradMethods.AUTO_DIFF,\n",
    "        eps=1e-2,\n",
    "    )(x, QuadCost(Q, p), dx)\n",
    "    \n",
    "    next_action = nominal_actions[0]\n",
    "    action_history.append(next_action)\n",
    "    u_init = torch.cat((nominal_actions[1:], torch.zeros(1, n_batch, dx.n_ctrl)), dim=0)\n",
    "    u_init[-2] = u_init[-3]\n",
    "    x = dx(x, next_action)\n",
    "\n",
    "    n_col = 4\n",
    "    n_row = n_batch // n_col\n",
    "    fig, axs = plt.subplots(n_row, n_col, figsize=(3*n_col,3*n_row), gridspec_kw = {'wspace':0, 'hspace':0})\n",
    "    axs = axs.reshape(-1)\n",
    "    for i in range(n_batch):\n",
    "        dx.get_frame(x[i], ax=axs[i])\n",
    "        axs[i].get_xaxis().set_visible(False)\n",
    "        axs[i].get_yaxis().set_visible(False)\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(os.path.join(t_dir, 'frame_{:03d}.png'.format(t)))\n",
    "    plt.close(fig)\n",
    "    \n",
    "action_history = torch.stack(action_history).detach()[:,:,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|████▊     | 48/100 [00:17<00:18,  2.83it/s]"
     ]
    }
   ],
   "source": [
    "# Plot actions\n",
    "for t in tqdm(range(T)):\n",
    "    fig, axs = plt.subplots(n_row, n_col, figsize=(3*n_col,3*n_row), gridspec_kw = {'wspace':0, 'hspace':0})\n",
    "    axs = axs.reshape(-1)\n",
    "    for i in range(n_batch):\n",
    "        axs[i].plot(action_history[:,i], color='k')\n",
    "        axs[i].set_ylim(-15, 15)\n",
    "        axs[i].axvline(t, color='k', ls='--', linewidth=4)\n",
    "        axs[i].get_xaxis().set_visible(False)\n",
    "        axs[i].get_yaxis().set_visible(False)\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(os.path.join(t_dir, 'actions_{:03d}.png'.format(t)))\n",
    "    plt.close(fig)\n",
    "    \n",
    "    f1 = os.path.join(t_dir, 'frame_{:03d}.png'.format(t))\n",
    "    f2 = os.path.join(t_dir, 'actions_{:03d}.png'.format(t))\n",
    "    f_out = os.path.join(t_dir, '{:03d}.png'.format(t))\n",
    "    os.system(f'convert {f1} {f2} +append -resize 1200x {f_out}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vid_fname = 'cartpole.mp4'\n",
    "\n",
    "if os.path.exists(vid_fname):\n",
    "    os.remove(vid_fname)\n",
    "    \n",
    "cmd = 'ffmpeg -r 16 -f image2 -i {}/%03d.png -vcodec libx264 -crf 25 -vf \"pad=ceil(iw/2)*2:ceil(ih/2)*2\" -pix_fmt yuv420p {}'.format(\n",
    "    t_dir, vid_fname\n",
    ")\n",
    "os.system(cmd)\n",
    "print('Saving video to: {}'.format(vid_fname))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "video = io.open(vid_fname, 'r+b').read()\n",
    "encoded = base64.b64encode(video)\n",
    "HTML(data='''<video alt=\"test\" controls>\n",
    "                <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" />\n",
    "             </video>'''.format(encoded.decode('ascii')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
